package com.webcat.dify.kb;

import com.alibaba.fastjson.JSON;
import com.webcat.dify.DifyProperties;
import com.webcat.dify.kb.entity.KnowledgeBase;
import com.webcat.dify.kb.entity.Page;
import com.webcat.dify.kb.entity.RequestQuery;
import jakarta.annotation.Resource;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

@Component
public class DifyKBClient {

    @Resource
    private DifyProperties difyProperties;

    @Resource
    private WebClient webClient;



    public Page getKnowledgeBase(int page, int limit) {
        String apiKey = difyProperties.getKbApiKey();
        String url = generateGetKBUrl(page, limit);
        return webClient.get().uri(url).header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .exchangeToMono(clientResponse -> clientResponse.bodyToMono(String.class)).map(data -> JSON.parseObject(data, Page.class))
                .map(p -> {
                    p.setData(JSON.parseArray(JSON.toJSONString(p.getData()), KnowledgeBase.class));
                    return p;
                }).block();
    }
    public Page getKnowledgeBase() {
        return getKnowledgeBase(1, 20);
    }


    public KBResponse getSyncKbResponse(String query, String datasetId) {
        RequestQuery requestQuery = new RequestQuery(query);
        String apiKey = difyProperties.getKbApiKey();
        return getSyncKbResponse(generateGetKbResponseUrl(datasetId), apiKey, requestQuery);
    }

    public Mono<KBResponse> getAsyncKbResponse(String query, String datasetId) {
        RequestQuery requestQuery = new RequestQuery(query);
        String apiKey = difyProperties.getKbApiKey();
        return getAsyncKbResponse(generateGetKbResponseUrl(datasetId), apiKey, requestQuery);
    }

    public KBResponse getSyncKbResponse(String url, String apiKey, RequestQuery requestQuery) {
        return webClient.post().uri(url).header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .bodyValue(requestQuery)
                .exchangeToMono(clientResponse -> clientResponse.bodyToMono(String.class)).map(data -> JSON.parseObject(data, KBResponse.class)).block();
    }

    public Mono<KBResponse> getAsyncKbResponse(String url, String apiKey, RequestQuery requestQuery) {
        return webClient.post().uri(url).header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
                .bodyValue(requestQuery)
                .exchangeToMono(clientResponse -> clientResponse.bodyToMono(String.class)).map(data -> JSON.parseObject(data, KBResponse.class));
    }

    /**
     * 生成检索知识库url
     * @param datasetId
     * @return
     */
    private String generateGetKbResponseUrl(String datasetId) {
        String baseUrl = difyProperties.getBaseUrl();
        String url;
        if (baseUrl.endsWith("/")) {
            url = baseUrl + "v1/datasets/" + datasetId + "/retrieve";
        }
        else {
            url = baseUrl + "/v1/datasets/" + datasetId + "/retrieve";
        }
        return url;
    }

    /**
     * 生成查询知识库列表url
     * @return
     */
    private String generateGetKBUrl(int page, int limit) {
        String baseUrl = difyProperties.getBaseUrl();
        String url;
        if (baseUrl.endsWith("/")) {
            url = baseUrl + "v1/datasets?page=" + page + "&limit=" + limit;
        }
        else {
            url = baseUrl + "/v1/datasets?page=" + page + "&limit=" + limit;
        }
        return url;
    }

    public static void main(String... args) {
        WebClient webClient = WebClient.builder()
                .build();


        Page page = webClient.get().uri("http://localhost/v1/datasets?page=1&limit=20").header(HttpHeaders.AUTHORIZATION, "Bearer dataset-Zw3oxGt300XTSvUXSivzZVRs")
                .exchangeToMono(clientResponse -> clientResponse.bodyToMono(String.class)).map(data -> JSON.parseObject(data, Page.class))
                .map(p -> {
                    p.setData(JSON.parseArray(JSON.toJSONString(p.getData()), KnowledgeBase.class));
                    return p;
                })
                .block();
        System.out.println(JSON.toJSONString(page));
    }
}
